import os
import dashscope
from dashscope import Generation
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
from dashscope import TextEmbedding


def getEmbedding(news):
    model = TextEmbedding.call(
        model=TextEmbedding.Models.text_embedding_v1,
        input=news
    )
    embeddings = [record['embedding'] for record in model.output['embeddings']]
    return embeddings if isinstance(news, list) else embeddings[0]

def getAnswer(query, context):
    prompt = f'''请基于```内的报道内容，回答我的问题。
	      ```
	      {context}
	      ```
	      我的问题是：{query}。
    '''

    rsp = Generation.call(model='qwen-turbo', prompt=prompt)
    return rsp.output.text


def search(text):
    # Search parameters for the index
    search_params = {
        "metric_type": "L2"
    }

    query_data = [-0.07788856327533722, 0.25216928124427795, -0.11765268445014954, -0.45044106245040894,
                  -0.43822047114372253, -0.30982694029808044, 0.2393917441368103, -0.5883796811103821,
                  0.2764524221420288, 0.2654382288455963, -0.2980782687664032, 0.04911184310913086, 0.21839120984077454,
                  0.01465060655027628, 0.5042624473571777, 0.7536708116531372, -0.18057098984718323,
                  -0.12001734226942062, -0.6089210510253906, 0.4411751329898834, 0.13465575873851776,
                  -0.3030729293823242, 0.13865312933921814, 0.03568962216377258, -0.4794985055923462,
                  -0.30423933267593384, 0.790494441986084, 0.058220550417900085, 0.30738455057144165,
                  0.3140547275543213, -0.10737757384777069, -0.1368374526500702, -0.1011720821261406,
                  -0.37301310896873474, 0.29350560903549194, -0.046796590089797974, -0.30051660537719727,
                  -0.21404476463794708, 0.011295430362224579, 0.1402393877506256, -0.14405351877212524,
                  -0.19671845436096191, 0.045492663979530334, 0.22737397253513336, 0.2547961473464966,
                  0.7625150084495544, 0.06825105845928192, -0.06948551535606384, 0.20536568760871887,
                  -0.08139297366142273, 0.5236148834228516, 3.7612662315368652, 0.8141233325004578, -0.25527423620224,
                  -0.0011854246258735657, 0.33001041412353516, 0.132666677236557, -0.5836617350578308,
                  -0.6639094948768616, 0.19722740352153778, -0.15364505350589752, -0.26286786794662476,
                  -0.34542733430862427, 0.014350846409797668, 0.3258463740348816, -0.10592664778232574,
                  -0.48532962799072266, 0.35756003856658936, 0.6863088607788086, -0.32747575640678406,
                  0.14576609432697296, -0.37993645668029785, 0.03933747112751007, 0.4857792854309082,
                  0.6312653422355652, 0.20663967728614807, 0.10216833651065826, 0.17953407764434814,
                  -0.07707604765892029, -0.5299369692802429, -0.2308180034160614, 0.33874452114105225,
                  -0.2688889503479004, -0.12979982793331146, -0.21032202243804932, 0.01876971125602722,
                  -0.03230294585227966, 0.3821732997894287, -0.6560024619102478, 0.24884629249572754,
                  -0.7004425525665283, 0.13797473907470703, 0.15133395791053772, -0.29248398542404175,
                  -0.310053288936615, -0.2763884961605072, 0.1538786143064499, 0.2335638850927353, 0.2913101315498352,
                  0.23923662304878235, -0.1037164032459259, -0.39591363072395325, 0.361303448677063,
                  0.36747393012046814, -0.41564512252807617, 0.09865429997444153, 0.9698308706283569,
                  -0.3892419934272766, -0.18374665081501007, -0.24489574134349823, -0.8977118134498596,
                  -0.33360105752944946, 0.2013644278049469, 0.3572555184364319, 0.4652222990989685, 0.1732686311006546,
                  0.5118612051010132, -0.19629818201065063, 0.2758632302284241, -0.2591860294342041,
                  0.05535241961479187, 0.028266653418540955, 0.3502984642982483, 0.120244100689888, 0.1552341878414154,
                  -0.17970529198646545, 0.270673543214798, -0.18956954777240753, -0.2925550043582916,
                  -0.03672260046005249, -0.19479605555534363, 0.12336871027946472, 0.5258193016052246,
                  -0.6438987255096436, -0.37592822313308716, 0.2935319244861603, -0.13084299862384796,
                  -0.1860196888446808, 0.13036027550697327, 1.0847046375274658, 0.06427684426307678,
                  -0.3635696470737457, -0.1818002462387085, -0.15886607766151428, -0.2999837100505829,
                  -0.04576520621776581, 0.11432917416095734, -0.014317028224468231, -0.21961726248264313,
                  -0.09239903092384338, -0.4757821261882782, -0.24967345595359802, -0.5079023838043213,
                  -0.2457471489906311, -0.029973991215229034, 0.19029366970062256, -0.3397367298603058,
                  0.19706469774246216, -0.26818081736564636, -0.039753664284944534, -0.19841879606246948,
                  0.04479853808879852, -0.4912832975387573, 0.11757799983024597, -0.026005566120147705,
                  -0.1456170380115509, 0.2583669126033783, -0.42830169200897217, -0.5070518851280212,
                  0.3022189140319824, 0.299455463886261, 0.1092911958694458, -0.11499342322349548, 0.4238845109939575,
                  0.07833890616893768, 0.09375515580177307, 0.44178926944732666, -0.1359623670578003,
                  -0.3766641914844513, -0.39801132678985596, 0.2955554723739624, -0.1142701581120491,
                  -0.2346135675907135, 0.03574724495410919, 0.07717245072126389, -0.34061571955680847,
                  -0.22885257005691528, -0.07690880447626114, 0.10522770881652832, -0.5659106373786926,
                  0.3138202130794525, -0.06570667773485184, 0.08993862569332123, 0.6058111190795898,
                  0.12517714500427246, -0.3181682229042053, -0.8022157549858093, -0.5426855087280273,
                  0.2900521159172058, 0.21444357931613922, -0.23010869324207306, -0.15728765726089478,
                  0.12447196245193481, 0.04778836667537689, -0.2874259054660797, 0.22052767872810364,
                  0.23176077008247375, 0.39928102493286133, 0.06598275899887085, 0.14622396230697632,
                  -0.9374401569366455, 0.3952939808368683, -0.4617876708507538, 0.30463266372680664, 0.420725017786026,
                  -0.2238750457763672, -0.9470816850662231, -0.341533899307251, 0.21379680931568146,
                  -0.1691327691078186, -0.013237284496426582, -0.5806558132171631, 0.11651848256587982,
                  0.20198720693588257, -0.014822624623775482, 0.1002756655216217, -0.1185879036784172,
                  -0.22719638049602509, -0.01651035249233246, -0.5543656349182129, -0.09866940975189209,
                  0.4142568111419678, -0.005775250494480133, -0.25561827421188354, -0.03287554532289505,
                  -0.15026149153709412, -0.16883602738380432, -0.09283998608589172, -0.16424909234046936,
                  0.2771233916282654, -0.2349776327610016, 0.07830783724784851, 0.4839682877063751,
                  -0.37361305952072144, -0.09133975207805634, 0.5906943678855896, -0.23859208822250366,
                  -0.6405794620513916, -0.015230894088745117, -0.4951673150062561, -0.2730823755264282,
                  0.6787944436073303, -0.12383382767438889, -0.0557321198284626, -0.3164328336715698,
                  0.18552199006080627, -0.05410326272249222, 0.06854148954153061, 0.5635229349136353,
                  0.7423043847084045, 0.16696110367774963, -0.1897403746843338, 1.004111886024475, -0.23876948654651642,
                  -0.24326670169830322, -0.5841752290725708, -0.23268964886665344, 0.278278648853302,
                  -0.005371566861867905, -0.019356831908226013, 0.20465491712093353, 0.19987794756889343,
                  -0.19475245475769043, -0.9052201509475708, -0.08986593782901764, 0.10055382549762726,
                  -0.2122848927974701, 0.08487393707036972, -0.6306875944137573, -0.5679085850715637,
                  -0.4084708094596863, 0.04024328291416168, 0.3558402359485626, 0.020627066493034363,
                  0.21997126936912537, 0.11393599212169647, -0.5834798216819763, 0.3143659830093384,
                  -0.5031486749649048, -0.27875617146492004, -0.03900480270385742, 0.9804702997207642,
                  -0.48834139108657837, -0.1993606984615326, 2.9212942123413086, -0.9602258801460266,
                  -0.6944621801376343, 0.2815970480442047, 0.5403702259063721, -0.7967433333396912, -0.0855545848608017,
                  0.45009735226631165, 0.010526403784751892, 0.10502420365810394, 0.050359878689050674,
                  -0.48477357625961304, 0.1065637618303299, -0.4814598858356476, 0.17914730310440063,
                  0.11324271559715271, -0.12431784719228745, 0.02805231511592865, 0.8864842653274536,
                  -0.2612030804157257, 0.8035739660263062, 0.2362087219953537, 0.1181323230266571, -0.6138833165168762,
                  -0.16065487265586853, -0.4513469338417053, 0.2138044536113739, 0.23681014776229858,
                  0.34282490611076355, -0.02923114225268364, 0.35761144757270813, 0.6654679775238037,
                  -0.3015890419483185, 0.387637734413147, -0.3302716016769409, -0.15886655449867249, 0.2591055631637573,
                  -0.4855079650878906, -0.3998069763183594, -0.3881744146347046, -0.061674103140830994,
                  -0.5670555830001831, 0.5555216073989868, -0.17824427783489227, 0.968442976474762, -0.3437103033065796,
                  0.5275223255157471, 0.4524834454059601, 0.5357840061187744, -0.3401309847831726, 0.24430297315120697,
                  0.17255164682865143, 0.4524039924144745, -0.5896444320678711, -0.4003605842590332,
                  -0.35884639620780945, 0.31397294998168945, 0.06253109872341156, -0.751962423324585,
                  -0.3858315944671631, 0.12643179297447205, -0.13133347034454346, -0.436269611120224,
                  0.36683493852615356, 0.5244601368904114, -0.26201438903808594, 0.3829519748687744,
                  -0.13638970255851746, -0.39315009117126465, -0.3367377519607544, 0.01931222341954708,
                  -0.4256916344165802, 0.2684188783168793, 0.5207522511482239, 0.1121944785118103, -0.7410165667533875,
                  -0.09967975318431854, -0.2019699215888977, -0.07970881462097168, 0.72531658411026,
                  0.11465223133563995, 0.08055596053600311, 0.04233785718679428, 0.38141757249832153,
                  0.08823555707931519, -0.08340170979499817, -0.2426895946264267, -0.07264497876167297,
                  0.13391999900341034, 1.0508573055267334, 0.1684572696685791, -0.5254690647125244, -0.2204318642616272,
                  -0.5850825905799866, -0.3808305263519287, 0.0692298412322998, 0.004826158285140991,
                  -0.28886809945106506, -0.14291173219680786, 0.1762162446975708, 0.4677971601486206,
                  0.43656182289123535, -0.19138862192630768, -0.03965768218040466, 0.04605989530682564,
                  0.11896360665559769, 0.41139891743659973, 0.43416348099708557, -0.19448035955429077,
                  0.37125730514526367, -0.42392146587371826, 0.43378570675849915, 0.15023919939994812,
                  0.10462220013141632, 0.6822272539138794, 0.1158829927444458, -0.3946756422519684, 0.6039702892303467,
                  0.08598807454109192, -0.624527096748352, 0.48408758640289307, -0.28135108947753906,
                  0.5355421900749207, 0.19042745232582092, 0.6921118497848511, -0.8088223934173584, 0.7858596444129944,
                  -0.36760208010673523, -0.8113353252410889, -0.13145124912261963, -0.2848861813545227,
                  0.21155884861946106, -0.0002547847107052803, 0.22506099939346313, -0.00960095226764679,
                  0.102645143866539, -0.3299243748188019, 0.010699950158596039, 0.1852046698331833, 0.19942982494831085,
                  -0.07719406485557556, -0.08679994195699692, 0.5051553249359131, 0.7059007287025452,
                  -0.3712652921676636, -0.101712167263031, 2.2917985916137695e-05, 0.4675922393798828,
                  0.13656623661518097, -0.5279770493507385, -0.3218238651752472, -0.2183079868555069,
                  0.25032147765159607, -0.6258039474487305, 0.0005197487771511078, 0.4076043367385864,
                  -0.3269498348236084, 0.2038053274154663, -0.48828405141830444, 0.47598105669021606,
                  -0.06750862300395966, 0.23539626598358154, -0.12597328424453735, -0.7738908529281616,
                  -0.05457894504070282, -0.1318793147802353, -0.15574859082698822, -0.27730199694633484,
                  0.2924756109714508, -0.10719689726829529, 0.42929142713546753, -0.17011545598506927,
                  -0.15078818798065186, -0.06611785292625427, -0.2660790681838989, -0.3147204518318176,
                  0.06074881553649902, 0.18515601754188538, 0.6222379803657532, 0.1661546528339386, 0.20109114050865173,
                  -0.2527843117713928, -0.0909147635102272, -0.0966176986694336, 0.13827574253082275,
                  -0.7537674903869629, -0.24310316145420074, 0.09009057283401489, -0.3780065178871155,
                  -0.3678011894226074, 0.36614546179771423, -0.2684903144836426, -0.20804932713508606,
                  -0.11938481777906418, 0.025689080357551575, -0.08850070834159851, -0.07870326936244965,
                  -0.06084885075688362, 0.011805176734924316, -0.3070593476295471, 0.303096204996109,
                  0.7757848501205444, -0.3838549256324768, 0.1720118224620819, -0.4687899947166443,
                  -0.40482592582702637, 0.6322567462921143, -0.002290278673171997, 0.18303993344306946,
                  -0.22217047214508057, 0.23760028183460236, -0.053253211081027985, -0.12301476299762726,
                  0.15726961195468903, 0.13854031264781952, -0.30954405665397644, 0.027304451912641525,
                  0.01829773560166359, -0.1851477175951004, 0.1977454423904419, 0.05948713421821594,
                  0.02271857112646103, 0.15461091697216034, 0.09718318283557892, -0.22596892714500427,
                  0.2728932499885559, 0.20680204033851624, -0.17937400937080383, 0.41079387068748474,
                  0.2658447027206421, 0.03152736276388168, 1.1198222637176514e-05, -0.3137202858924866,
                  0.334148645401001, -0.5278273820877075, -0.10891255736351013, -0.4482111930847168,
                  -0.6056948304176331, 0.1143798679113388, 0.3899139165878296, -0.4928843379020691, 0.19867286086082458,
                  -0.13204844295978546, -0.16546683013439178, 0.020817190408706665, -0.4772500991821289,
                  -0.44608166813850403, 0.4983307719230652, 0.11634549498558044, 0.4837114214897156,
                  0.25577399134635925, 0.20483030378818512, -0.08691994845867157, -0.019514024257659912,
                  -0.6728929877281189, -0.18899255990982056, -1.043459415435791, -0.24438245594501495,
                  0.17328165471553802, -0.18468081951141357, 0.1613994836807251, -0.0017190277576446533,
                  0.2335360050201416, 0.4598412811756134, 0.7197536826133728, 0.03669536113739014, 0.11185277998447418,
                  0.09560612589120865, -0.4345240592956543, -0.4037734270095825, -0.7996190190315247,
                  -0.1579737812280655, -0.11487022042274475, -0.011134207248687744, -0.10684458166360855,
                  -0.020331289619207382, -0.11303995549678802, -1.2621780633926392, -0.2792239487171173,
                  0.33870336413383484, -0.026915714144706726, 0.037878263741731644, 0.4874606132507324,
                  -0.2564655542373657, -0.2525896728038788, -0.8261376619338989, -0.21470916271209717,
                  -0.07947099953889847, -0.11993465572595596, -0.691865861415863, -0.5216588973999023,
                  -0.7579572796821594, 0.6087493300437927, -0.6785551309585571, -0.17529962956905365,
                  -0.2942824363708496, -0.42323043942451477, 0.22854886949062347, -0.8192356824874878,
                  0.034588560461997986, -0.15592025220394135, -0.35614925622940063, 0.025347143411636353,
                  0.2592090368270874, 0.007428936660289764, -0.02043507993221283, -0.05922795087099075,
                  -0.3018490970134735, 0.7524340152740479, -0.22349926829338074, -0.09875532984733582,
                  -1.1669083833694458, 0.3223028779029846, -0.045816466212272644, -0.5933247208595276,
                  -0.7008919715881348, -0.21919968724250793, 0.4993739128112793, -0.15214313566684723,
                  0.5298256278038025, 0.26494428515434265, -0.4337862432003021, -0.2241154909133911,
                  -0.1877327561378479, 0.09065397083759308, 0.0005276650190353394, 0.8502806425094604,
                  -0.4867962598800659, 0.3032952547073364, 0.06367290765047073, 0.0306943878531456, 0.4113042950630188,
                  0.162433922290802, 0.04151131212711334, -0.2990403473377228, -0.21312367916107178, 0.8790697455406189,
                  -0.18938425183296204, 0.05469211935997009, -0.37049350142478943, 0.0004969090223312378,
                  0.20864440500736237, -0.2165592461824417, 0.152040034532547, -0.1483556628227234, 0.6677651405334473,
                  0.029458023607730865, 0.0793645977973938, -0.2202182412147522, -0.06631152331829071,
                  -0.33550697565078735, 0.0505312904715538, -0.15904739499092102, 0.018997633829712868,
                  0.062014833092689514, 0.36965805292129517, -0.02521061897277832, 0.22892868518829346,
                  0.058896031230688095, 0.6671242117881775, 0.10001083463430405, 0.22409012913703918,
                  -0.2078913450241089, -0.1857926845550537, -0.3200269341468811, -0.3500043749809265,
                  0.06568434834480286, 0.19164150953292847, 0.5786838531494141, 0.0871022641658783, -0.1439243108034134,
                  -0.22696168720722198, 0.013963282108306885, -0.1761934608221054, -0.1668201982975006,
                  0.4594777822494507, 0.2123977243900299, -0.5250391364097595, 0.5836950540542603, -0.5521676540374756,
                  0.6342851519584656, -0.2104732096195221, 0.34964898228645325, -0.24696463346481323,
                  0.0434231162071228, 0.4729090929031372, -0.6382764577865601, 0.15553952753543854, 0.39836376905441284,
                  0.10604807734489441, 0.03988282382488251, 0.362132728099823, -0.3800339102745056, -0.351187139749527,
                  -0.35104021430015564, 0.2087949514389038, -0.29671385884284973, 0.8005974292755127,
                  0.2990625202655792, 0.1797705888748169, 0.26352912187576294, 1.0752153396606445, -0.4182196259498596,
                  0.08031558990478516, 0.2519063949584961, -0.023489098995923996, 0.08062012493610382,
                  0.5145708918571472, -0.033798910677433014, -0.7347238659858704, -0.5857511758804321,
                  0.07898667454719543, -0.6800729036331177, 0.8050336241722107, -0.18232931196689606,
                  -0.15275052189826965, 0.23682239651679993, -0.936772346496582, -0.017802417278289795,
                  -0.12631836533546448, 0.1303652822971344, -0.2023494839668274, 0.48870450258255005,
                  0.3178473711013794, -0.1006137877702713, 0.17500999569892883, 0.3160974979400635,
                  -0.014872327446937561, -0.08070410788059235, -0.06064004451036453, 0.5445604920387268,
                  0.5043504238128662, -0.3738894760608673, 0.16712115705013275, 0.0163232684135437, 0.22907933592796326,
                  -0.25745218992233276, 0.12479446828365326, -0.2728506922721863, 0.01609943062067032,
                  -0.2711203992366791, -0.2737160325050354, -0.09899036586284637, 0.09446194767951965,
                  0.5107632875442505, -0.026948586106300354, 0.004189789295196533, -0.26581576466560364,
                  0.3054140508174896, -0.3571496903896332, 0.11436869949102402, 0.23725557327270508,
                  0.11356263607740402, -0.24781283736228943, 0.10413742065429688, 0.14235469698905945,
                  -0.03406541794538498, 0.284166544675827, -0.18688184022903442, -0.14589504897594452,
                  -0.3320857286453247, 0.012831784784793854, -0.9738390445709229, 0.4511360824108124,
                  -0.3725200891494751, -0.5290850400924683, -0.30410289764404297, 0.6612716317176819,
                  -0.2758220136165619, -0.023134391754865646, -0.08407595008611679, -0.31962937116622925,
                  -0.19127896428108215, -0.24641035497188568]
    results = collection.search(
        data=[query_data],  # Embeded search value
        anns_field="embedding",  # Search across embeddings
        param=search_params,
        limit=20,  # Limit to five results per search
        output_fields=['text']  # Include title field in result
    )

    ret = []
    for hit in results[0]:
        ret.append(hit.entity.get('text'))
    return ret


if __name__ == '__main__':

    current_path = os.path.abspath(os.path.dirname(__file__))   # 当前目录
    root_path = os.path.abspath(os.path.join(current_path, '..'))   # 上级目录
    data_path = f'{root_path}/CEC-Corpus/raw corpus/allSourceText'

    # 配置Dashscope API KEY
    dashscope.api_key = '<YOUR_DASHSCOPE_API_KEY>'

    # 配置Milvus参数
    COLLECTION_NAME = 'CEC_Corpus'
    DIMENSION = 768
    MILVUS_HOST = 'app1'
    MILVUS_PORT = '19530'
    USER = 'root'
    PASSWORD = '<password>'

    connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)

    fields = [
        FieldSchema(name='id', dtype=DataType.INT64, descrition='Ids', is_primary=True, auto_id=True),
        FieldSchema(name='text', dtype=DataType.VARCHAR, description='Text', max_length=4096),
        FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, description='Embedding vectors', dim=DIMENSION)
    ]
    schema = CollectionSchema(fields=fields, description='CEC Corpus Collection')
    collection = Collection(name=COLLECTION_NAME,schema=schema)

    # Load the collection into memory for searching
    collection.load()

    question = '北京中央电视台工地发生大火，发生在哪里？出动了多少辆消防车？人员伤亡情况如何？'
    context = search(question)
    print(context)
    # answer = getAnswer(question, context)
    # print(answer)